package com.example.iot.service.websocket.impl;

import com.baomidou.mybatisplus.core.toolkit.CollectionUtils;
import com.example.iot.entity.dto.TsData;
import com.example.iot.entity.id.EntityId;
import com.example.iot.service.websocket.SubscriptionService;
import com.example.iot.ts.TimeSeriesSubscriptionInfo;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;

import javax.annotation.PostConstruct;
import javax.annotation.PreDestroy;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

/**
 * @author ChenXiangpeng
 */
@Service
@Slf4j
public class DefaultSubscriptionServiceImpl implements SubscriptionService {
    /**
     * sesseionId -> Session
     */
    private final Map<String, TimeSeriesSubscriptionInfo<String>> subscriptionsBySessionId = new ConcurrentHashMap<>();
    /**
     * entityId -> Session
     */
    private final Map<EntityId, Set<TimeSeriesSubscriptionInfo<String>>> subscriptionsByEntityId = new ConcurrentHashMap<>();

    private final ObjectMapper objectMapper = new ObjectMapper();
    private ExecutorService subscriptionUpdateExecutor;

    @PostConstruct
    public void initExecutor() {
        subscriptionUpdateExecutor = Executors.newWorkStealingPool(20);
    }

    @PreDestroy
    public void shutdownExecutor() {
        if (subscriptionUpdateExecutor != null) {
            subscriptionUpdateExecutor.shutdownNow();
        }
    }

    @Override
    public void addSubscription(TimeSeriesSubscriptionInfo<String> subscriptionInfo) {
        subscriptionsBySessionId.putIfAbsent(subscriptionInfo.getSessionId(), subscriptionInfo);
        subscriptionsByEntityId
                .computeIfAbsent(subscriptionInfo.getEntityId(), k -> ConcurrentHashMap.newKeySet()).add(subscriptionInfo);
    }

    // todo 方法中可以添加自定义Consumer的过滤以及处理方法
    @Override
    public void onTimeSeriesUpdate(EntityId entityId, TsData tsData) {
        Set<TimeSeriesSubscriptionInfo<String>> timeSeriesSubscriptionInfos = subscriptionsByEntityId.getOrDefault(entityId, ConcurrentHashMap.newKeySet());
        // 对该实体的所有Session进行循环消息推送
        timeSeriesSubscriptionInfos.forEach(e->{
            try {
                e.getUpdateConsumer().accept(e.getSessionId(), objectMapper.writeValueAsString(tsData));
            } catch (JsonProcessingException jsonProcessingException) {
                jsonProcessingException.printStackTrace();
            }
        });
    }

    @Override
    public void cancelAllSubscription(String sessionId) {
        TimeSeriesSubscriptionInfo<String> remove = subscriptionsBySessionId.remove(sessionId);
        Set<TimeSeriesSubscriptionInfo<String>> timeSeriesSubscriptionInfos = subscriptionsByEntityId.get(remove.getEntityId());
        if (CollectionUtils.isNotEmpty(timeSeriesSubscriptionInfos)){
            timeSeriesSubscriptionInfos.remove(remove);
        }
    }
}
